import librosa
from librosa import display
import matplotlib.pyplot as plt
import numpy as np
from pydub import AudioSegment
import os
AUDIO_DIR = "data/genres_original"
MEL_SPECTROGRAM_PATH = "data/mel_spectrogram"
SPLIT_SONG_DIR = "data/content"

if not os.path.exists(MEL_SPECTROGRAM_PATH):
    os.mkdir(MEL_SPECTROGRAM_PATH, 0o755)

genres = "blues classical country disco hiphop jazz metal pop reggae rock".split()
# genres = "metal pop reggae rock".split()
#


def audio_to_mel_spectrogram(genre, audio_file):
    save_path = os.path.join(
        os.path.join(MEL_SPECTROGRAM_PATH, genre, audio_file.replace(".", "").replace("wav", ".png")))

    audio_recording = os.path.join(SPLIT_SONG_DIR, genre, audio_file)
    data, sr = librosa.load(audio_recording)
    sgram = librosa.stft(data)
    sgram_mag, _ = librosa.magphase(sgram)
    mel_scale_sgram = librosa.feature.melspectrogram(S=sgram_mag, sr=sr)
    mel_sgram = librosa.amplitude_to_db(mel_scale_sgram, ref=np.min)
    librosa.display.specshow(mel_sgram, sr=sr)

    plt.savefig(save_path, bbox_inches='tight', pad_inches=-0.045, dpi=128)

def split_song(genre, song):
    audio = os.path.join(AUDIO_DIR, genre, song)
    audio = AudioSegment.from_wav(audio)
    num = int(song.split('.')[1])
    for i in range(10):
        t1 = 3 * i * 1000
        t2 = 3 * (i + 1) * 1000
        if i == 9:
            new_audio = audio[t1:]
        else:
            new_audio = audio[t1:t2]
        save_path = os.path.join(SPLIT_SONG_DIR, genre, genre + "_{:05d}.wav".format(num * 10 + i))
        new_audio.export(save_path, format='wav')


# for genre in genres:
#     songs = os.listdir(os.path.join(AUDIO_DIR, genre))
#     for song in songs:
#         split_song(genre, song)

# for genre in genres:
#     songs = os.listdir(os.path.join(SPLIT_SONG_DIR, genre))
#     print(genre, len(songs))
#
#
for genre in genres:
    if not os.path.exists(MEL_SPECTROGRAM_PATH + "/" + genre):
        os.mkdir(MEL_SPECTROGRAM_PATH + "/" + genre, 0o755)
    genre_path = os.path.join(SPLIT_SONG_DIR, genre)
    for audio in os.listdir(genre_path):
        try:
            audio_to_mel_spectrogram(genre, audio)
        except Exception:
            print(genre, audio)
            continue

